#include "TcpServer.h"
#include <algorithm>

using namespace clusnet::net;

TcpServer::TcpServer(const char* host, unsigned short port)
    : TcpServer(boost::asio::ip::tcp::endpoint(boost::asio::ip::make_address(host), port))
{
}

TcpServer::TcpServer(unsigned short port, const boost::asio::ip::tcp::resolver::protocol_type& proto)
    : TcpServer(boost::asio::ip::tcp::endpoint(proto, port))
{
}

TcpServer::TcpServer(const boost::asio::ip::tcp::endpoint& endpoint)
    : acceptor_(this, mainIO_, endpoint),
      next_(0),
      newConnCallback_(nullptr),
      closeCallback_(nullptr),
      writeCallback_(nullptr),
      messageCallback_(nullptr)
{
    struct sigaction sa;
    sa.sa_handler = SIG_IGN;
    sigaction(SIGPIPE, &sa, nullptr);

    acceptor_.setNewConnectionCallback(std::bind(&TcpServer::handleNewConnection, this, std::placeholders::_1));
    acceptor_.setCloseCallback(std::bind(&TcpServer::handleCloseConnection, this, std::placeholders::_1));
    acceptor_.setWriteCompleteCallback(std::bind(&TcpServer::handleWriteComplete, this, std::placeholders::_1, std::placeholders::_2));
    acceptor_.setMessageCallback(std::bind(&TcpServer::handleMessage, this, std::placeholders::_1, std::placeholders::_2));
}

TcpServer::~TcpServer()
{
    conns_.clear();
    assert(mthreads_.empty());
}

// 1、only one io_context run in multi threads
// 2、multi io_contexts run in multi threads, every thread has its own io_context
void TcpServer::start(size_t numThreads, size_t ioNum)
{
    assert(wthreads_.empty() && mthreads_.empty() && workIOs_.empty() && workers_.empty() && (numThreads > 0 || ioNum > 0));

    if (ioNum > 0)
    {
        numThreads = ioNum + 1;
        std::vector<boost::asio::io_context>    ios(ioNum);
        workIOs_.swap(ios);
    }

    acceptor_.safeStartAccept();

    wthreads_.reserve(ioNum);
    mthreads_.reserve(numThreads - ioNum);
    workers_.reserve(ioNum);

    for (size_t i = 0; i != ioNum; ++i)
    {
        workers_.emplace_back(boost::asio::make_work_guard(workIOs_[i]));
        std::string name = "tcpserver_w" + std::to_string(i + 1);
        wthreads_.emplace_back(new clusnet::base::Thread([this, i](){ workIOs_[i].run(); }, name));
        (*wthreads_.rbegin())->start();
    }

    for (size_t i = ioNum; i != numThreads; ++i)
    {
        std::string name = "tcpserver_m" + std::to_string(i - ioNum + 1);
        mthreads_.emplace_back(new clusnet::base::Thread([this](){ mainIO_.run(); }, name));
        (*mthreads_.rbegin())->start();
    }

    LOG(INFO) << "tcp server start ok. " << numThreads << " threads " << "and " << ioNum + 1 << " io_context";
}

void TcpServer::stop()
{
    assert(!mthreads_.empty());
    auto work = boost::asio::make_work_guard(mainIO_);
    acceptor_.safeStop([this]{
        std::for_each(conns_.begin(), conns_.end(), [](const TcpConnectionPtr& conn) { conn->safeShutDown(); });
    });
    std::for_each(workers_.begin(), workers_.end(), [](Work& work) { work.reset(); });

    // let workIO thread quit at first, then mainIO
    std::for_each(wthreads_.begin(), wthreads_.end(), std::bind(&clusnet::base::Thread::join, std::placeholders::_1));
    wthreads_.clear();

    work.reset();
    std::for_each(mthreads_.begin(), mthreads_.end(), std::bind(&clusnet::base::Thread::join, std::placeholders::_1));
    mthreads_.clear();
}

void TcpServer::handleNewConnection(TcpConnectionPtr& conn)
{
    conns_.emplace(conn);
    if (newConnCallback_)
    {
        newConnCallback_(conn);
    }
    else
    {
        conn->safeStartReceive();
    }
}

void TcpServer::handleCloseConnection(TcpConnectionPtr& conn)
{
    if (closeCallback_)
    {
        closeCallback_(conn);
    }
    conns_.erase(conn);
}

void TcpServer::handleWriteComplete(TcpConnectionPtr& conn, const TcpDataPacketPtr& data)
{
    if (writeCallback_)
    {
        writeCallback_(conn, data);
    }
}

void TcpServer::handleMessage(TcpConnectionPtr& conn, TcpUniqDataPacketPtr&& data)
{
    assert(messageCallback_ != nullptr);
    messageCallback_(conn, std::move(data));
}

boost::asio::io_context& TcpServer::getNextIoContext()
{
    if (!workIOs_.empty())
    {
        if (next_ >= workIOs_.size())
        {
            next_ = 0;
        }
        return workIOs_[next_++];
    }

    return mainIO_;
}
